#include <fstream>
#include <iostream>
#include <math.h>
#include <sstream>
#include <string>
#include <vector>
using namespace std;

#include "Bitext.hh"
#include "Model.hh"
#include "Tokenizer.hh"

void SentencePair::print(Bitext &bitext, ostream &out){
  // print pair
  out << source << endl;
  out << target << endl;

  // print labeling 
  out << "{";
  for(set<pair<int, int> >::iterator itr = alignment_labeling.begin(); itr != alignment_labeling.end(); itr++){
    int j = itr->first;
    int i = itr->second;
    string token_source = source_tokenizer[j];
    string token_target = target_tokenizer[i];
    out << " [" << j << " (" << token_source << ") -> " << i << " (" << token_target << ") ]";
  }
  out << " }" << endl;
}

void SentencePair::print_aachen(Bitext &bitext, int c, double p, ostream &out){
  //# Sentence pair (1) source length 3 target length 3 alignment score : 0.0566893
  //america and china 
  //NULL ({ 2 3 }) amerika ({ 1 }) und ({ }) china ({ }) 

  int J = source_ids.size();
  int I = target_ids.size();

  // print general info
  out << "# Sentence pair (" << c << ") source length " << J
      << " target length " << I
      << " alignment score " << p << endl;

  // source sentence
  for(int j = 0; j < J; j++){
    out << bitext.source_vocabulary.map_to_token(source_ids[j]) << " ";
  }
  out << endl;

  // print viterbi alignment
  for(int i = 0; i < I; i++){
    out << bitext.target_vocabulary.map_to_token(target_ids[i]) << " ({ ";
    for(int j = 0; j < J; j++){
      if(alignment.find(pair<int, int>(j, i)) != alignment.end()){
	out << j + 1 << " ";
      }
    }
    out << "}) ";
  }
  out << endl;
}

bool SentencePair::is_labeled_training_example(){
  return labeled && for_train;
}

bool SentencePair::is_labeled_testing_example(){
  return labeled && !for_train;
}

Bitext::Bitext(string fn) : ifs(fn.c_str()){
  read_so_far = 0;
  filename = fn;

  // check that could open bitext...
  if(!ifs.is_open()){
    cout << "could not open bitext!" << endl;
    exit(1);
  }
}

string Bitext::get_filename(){
  return filename;
}

void Bitext::inspect(Model *p_model, string summary_filename, string verbose_filename, string viterbi_filename, int iteration){
  // get ready for verbose output
  stringstream i_s;
  i_s << iteration;
  string verbose_dump_filename = verbose_filename + ".verbose." + p_model->get_model_name() + ".i" + i_s.str();
  ofstream verbose;
  verbose.open(verbose_dump_filename.c_str());
  cout << "printing verbose output to " << verbose_dump_filename << endl;

  // get ready for viterbi output
  string viterbi_dump_filename = viterbi_filename + ".viterbi." + p_model->get_model_name() + ".i" + i_s.str();
  ofstream viterbi;
  viterbi.open(viterbi_dump_filename.c_str());
  cout << "printing viterbi alignment to " << viterbi_dump_filename << endl;

  // start inspecting model over dataset..
  rewind();
  SentencePair pair;
  double sum_aer = 0.0;
  double sum_ssaer = 0.0;
  int count = 0;
  double perplexity_viterbi = 0.0;
  int viterbi_smoothed = 0;
  while(next(pair)){
    // viterbi output
    pair.alignment.clear();
    double p_viterbi = p_model->viterbi(pair, true);
    pair.print_aachen(*this, count + 1, p_viterbi, viterbi);
    if(p_viterbi == 0.0){
      viterbi_smoothed++;
      p_viterbi += 10e-10;
    }
    perplexity_viterbi -= p_viterbi * log2(p_viterbi);

    if(pair.is_labeled_testing_example()){
      // print to verbose output: pair, golden labeling, viterbi prob
      pair.print(*this, verbose);
      verbose << p_viterbi << endl;

      // calculate AER
      set<std::pair<int, int> > a_union;
      set_union(pair.alignment_labeling.begin(), pair.alignment_labeling.end(),
		pair.alignment.begin(),          pair.alignment.end(),
		inserter(a_union, a_union.end()));
      set<std::pair<int, int> > a_intersection;
      set_intersection(pair.alignment_labeling.begin(), pair.alignment_labeling.end(),
		       pair.alignment.begin(),          pair.alignment.end(),
		       inserter(a_intersection, a_intersection.end()));
      double aer = 1.0 - (a_intersection.size() / (double) a_union.size());
      verbose << "AER   = " << aer << endl;
      sum_aer += aer;

      // calculate ssAER: single-source AER, where only best-scoring mapping of source->target is assessed
      set<std::pair<int, int> > a_modified_intersection = a_intersection;
      set<std::pair<int, int> > a_modified_union = a_union;
      for(set<std::pair<int, int> >::iterator itr = a_intersection.begin(); itr != a_intersection.end(); itr++){
	int j = itr->first;
	int i = itr->second;
	// delete other (j',*) alignments from intersection and union.
	for(set<std::pair<int, int> >::iterator itr2 = a_union.begin(); itr2 != a_union.end(); itr2++){
	  int j2 = itr2->first;
	  int i2 = itr2->second;
	  if((j2 == j) && (i2 != i)){
	    if(a_modified_intersection.find(*itr) != a_modified_intersection.end())
	      a_modified_intersection.erase(*itr);
	    a_modified_union.erase(*itr);
	  }
	}
      }
      // all remaining one->many alignments in union .. leave only one. (we can just count this.)
      set<int> covered_source;
      for(set<std::pair<int, int> >::iterator itr = a_modified_union.begin(); itr != a_modified_union.end(); itr++){
	covered_source.insert(itr->first);
      }
      double ssaer = 1.0 - (a_modified_intersection.size() / (double) covered_source.size());
      verbose << "ssAER = " << ssaer << endl;
      sum_ssaer += ssaer;

      // for divisor in average
      count++;
    }
  }

  if(viterbi_smoothed > 0){
    cout << "warning: " << viterbi_smoothed << " viterbi probabilities smoothed for the sake of perplexity calculation" << endl;
  }

  // output perplexity of viterbi alignment
  cout << "viterbi perplexity = " << pow(2.0, perplexity_viterbi) << endl;

  // find mean of metrics
  double mean_aer = sum_aer / (double) count;
  double mean_ssaer = sum_ssaer / (double) count;

  // print summary to file ..
  stringstream i_string;
  i_string << iteration;
  string summary_dump_filename = summary_filename + ".summary." + p_model->get_model_name() + ".i" + i_string.str();
  cout << "printing summary to " << summary_dump_filename << endl;
  ofstream summary;
  summary.open(summary_dump_filename.c_str());
  summary << "mean AER   = " << mean_aer << endl;
  summary << "mean ssAER = " << mean_ssaer << endl;
  summary.close();

  // close up verbose output.
  verbose.close();
  viterbi.close();
}

bool Bitext::next(SentencePair &pair){  
  // read pair
  string s;
  if(!getline(ifs, s)){
    cerr << "(" << read_so_far << ") ";
    return false;
  }
  stringstream ss(s);
  ss >> pair.n;
  if(!getline(ifs, pair.source)){
    cerr << "got n, but no source!" << endl;
    exit(1);
  }
  if(!getline(ifs, pair.target)){
    cerr << "got n and source, but no target!" << endl;
    exit(1);
  }
  string align;
  if(!getline(ifs, align)){
    cerr << "got n, source, and target, but no alignment!" << endl;
    exit(1);
  }

  // tokenize source and target & add null word
  pair.source_tokenizer.tokenize(pair.source);
  pair.target = "NULL " + pair.target;
  pair.target_tokenizer.tokenize(pair.target);
  if((pair.J = pair.source_tokenizer.size()) == 0){
    cerr << "warning: pair " << read_so_far << " has zero source length, skipping" << endl;
    return next(pair);
  }
  pair.I = pair.target_tokenizer.size();

  // add to vocabulary
  source_vocabulary.add_to_vocabulary(pair.source_tokenizer);
  target_vocabulary.add_to_vocabulary(pair.target_tokenizer);

  // convert tokens to ids
  source_vocabulary.map_to_ids(pair.source_tokenizer, pair.source_ids);
  target_vocabulary.map_to_ids(pair.target_tokenizer, pair.target_ids);

  // read in alignment source->target map
  pair.alignment_labeling.clear();
  pair.labeled = (align != "%");
  if(pair.labeled){
    stringstream ss_align(align);
    string train_or_test;
    ss_align >> train_or_test;
    if(train_or_test == "train"){
      pair.for_train = true;
    }else if(train_or_test == "test"){
      pair.for_train = false;
    }else{
      cerr << "got \"" << train_or_test << "\", should be \"train\" or \"test\"" << endl;
      exit(1);
    }
    int j;
    while(ss_align >> j){
      int i;
      if(!(ss_align >> i)){
      cerr << "lopsided alignment!" << endl;
      exit(1);
      }
      if(ss.fail()){
	cerr << "bad alignment!" << endl;
      }
      pair.alignment_labeling.insert(std::pair<int, int>(j, i + 1));
    }
  }

  // output .. ?
  if(++read_so_far % 10000 == 0){
    cerr << "(" << read_so_far << ") ";
  }

  return true;
}

void Bitext::rewind(){
  read_so_far = 0;

  ifs.clear();
  ifs.seekg(0, ios::beg);
}
